import numpy as np


def kmeans_search(X, y, target):
    def euclidean_distances(x, y):
        return np.sqrt(np.sum((x - y) ** 2))

    print("="*50)
    print("\t\t\t\t进入相似患者搜索系统")
    print("="*50)
    print("目标患者：", target)

    # 搜索与目标患者最相似的样本，作为第一个初始质心，固定
    mostlike = (None, np.inf)
    for i in range(len(X)):
        if euclidean_distances(target, X[i]) < mostlike[1]:
            mostlike = (i, euclidean_distances(target, X[i]))

    # 距离第一个初始质心最远的样本作为第二个初始质心，随着算法迭代而不断更新
    mostunlike = (None, 0)
    for i in range(len(X)):
        if euclidean_distances(mostlike[0], X[i]) > mostunlike[1]:
            mostunlike = (i, euclidean_distances(mostlike[0], X[i]))

    # 初始质心
    cores = [X[mostlike[0]], X[mostunlike[0]]]
    print("初始质心：", cores)

    # 迭代，直到质心不再改变
    print("="*50)
    print("\t\t\t\t\t开始聚类")
    counter = 1
    while True:
        print(f"第{counter}次迭代...")
        counter += 1
        l1, l2 = [], []
        labels = []
        for i in range(len(X)):
            if euclidean_distances(X[i], cores[0]) < euclidean_distances(X[i], cores[1]):
                l1.append(X[i])
                labels.append(i)
            else:
                l2.append((X[i]))
        pre = cores[1]
        cores[1] = sum(l2) / len(l2)  # 更新质心
        print("\t更新质心：", pre, "->", cores[1])
        if (pre - cores[1]).sum() == 0 or counter >= 1000:
            print("\t\t\t\t\t搜索完毕")
            print("=" * 50)
            return np.array(l1), y[labels]


if __name__ == '__main__':
    X = np.array(range(100)).reshape(4, 25)
    y = np.array([0, 1, 1, 0])
    target = np.array([[3, 2, 8, 5, 3, 3, 2, 8, 5, 3, 3, 2, 8, 5, 3, 3, 2, 8, 5, 3, 3, 2, 8, 5, 3]])
    X_train, y_train = kmeans_search(X, y, target)
